# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
# generator_ascend_generate_grouped_matmul_add.py
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul")
class GroupedMatmulAddGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)
        self.m = random.randint(1,2048)
        self.k = random.randint(1,7168)
        self.n = random.randint(1,7168)
        self.e = 0


    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        x_shape = case_config.inputs[0][0].shape
        w_shape = case_config.inputs[1][0].shape
        bias_shape = case_config.inputs[2][0].shape
        antiquant_scale_shape = case_config.inputs[3][0].shape
        antiquant_offset_shape = case_config.inputs[4][0].shape
        group_list_shape = case_config.inputs[5].shape
        
        # x shape
        x_shape[0] = self.m  # m
        self.e = random.randint(1, min(x_shape[0], 128))
        x_shape[1] = self.k # k
        
        # weight shape
        w_shape[0] = self.e # e
        w_shape[1] = self.k # k
        w_shape[2] = self.n # n

        # bias shape
        bias_shape[0] = self.e
        bias_shape[1] = self.n 

        # antiquant scale shape
        antiquant_scale_shape[0] = self.e
        antiquant_scale_shape[1] = self.n
    
        # antiquant offset shape
        antiquant_offset_shape[0] = self.e
        antiquant_offset_shape[1] = self.n

        #groupList 
        group_list_shape[0] = self.e
        case_config.inputs[5].range_values = [1, self.m]

        return case_config
    